import os
import glob
from tensorboard.backend.event_processing import event_accumulator
import matplotlib.pyplot as plt
import numpy as np

def find_tensorboard_logs(base_dir="f:/curWork/light_mappo/results"):
    """查找所有TensorBoard日志文件"""
    log_dirs = []
    for root, dirs, files in os.walk(base_dir):
        for file in files:
            if file.startswith("events.out.tfevents"):
                log_dirs.append(root)
                break
    return log_dirs

def analyze_tensorboard_logs(log_dir):
    """分析TensorBoard日志文件"""
    print(f"分析日志目录: {log_dir}")
    ea = event_accumulator.EventAccumulator(log_dir, size_guidance={
        event_accumulator.SCALARS: 0,  # 加载所有标量
    })
    ea.Reload()
    
    # 获取可用的标量标签
    tags = ea.Tags()['scalars']
    print(f"可用的指标: {tags}")
    
    if not tags:
        print("未找到任何指标，请检查日志路径是否正确")
        return None
    
    # 收集所有指标数据
    metrics_data = {}
    for tag in tags:
        events = ea.Scalars(tag)
        metrics_data[tag] = {
            'steps': [e.step for e in events],
            'values': [e.value for e in events],
            'wall_time': [e.wall_time for e in events]
        }
        
        # 打印指标统计信息
        values = metrics_data[tag]['values']
        if values:
            print(f"\n指标: {tag}")
            print(f"  数据点数量: {len(values)}")
            print(f"  平均值: {np.mean(values):.4f}")
            print(f"  最大值: {np.max(values):.4f}")
            print(f"  最小值: {np.min(values):.4f}")
            print(f"  最终值: {values[-1]:.4f}")
    
    return metrics_data

def analyze_run_directory(run_dir):
    """分析整个运行目录及其所有子目录"""
    print(f"分析运行目录: {run_dir}")
    
    # 查找该运行目录下的所有日志目录
    log_dirs = []
    for root, dirs, files in os.walk(run_dir):
        for file in files:
            if file.startswith("events.out.tfevents"):
                log_dirs.append(root)
                break
    
    if not log_dirs:
        print("未找到任何TensorBoard日志文件")
        return
    
    print(f"找到 {len(log_dirs)} 个日志子目录")
    
    # 收集所有指标数据
    all_metrics = {}
    
    for log_dir in log_dirs:
        # 提取相对路径作为指标前缀
        rel_path = os.path.relpath(log_dir, run_dir)
        prefix = rel_path.replace("\\", "/")
        
        # 分析日志
        metrics_data = analyze_tensorboard_logs(log_dir)
        
        if metrics_data:
            # 添加前缀并合并到总指标中
            for tag, data in metrics_data.items():
                full_tag = f"{prefix}/{tag}" if prefix != "." else tag
                all_metrics[full_tag] = data
    
    # 绘制所有指标
    if all_metrics:
        save_dir = os.path.join(run_dir, "analysis")
        plot_metrics(all_metrics, save_dir)
        print(f"\n图表已保存到: {save_dir}")
    else:
        print("未找到任何有效指标数据")

def plot_metrics(metrics_data, save_dir=None):
    """绘制指标图表"""
    if not metrics_data:
        return
    
    # 创建保存目录
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
    
    # 按指标类型分组
    metric_types = {}
    for full_tag, data in metrics_data.items():
        # 提取基本指标类型（例如从"agent0/value_loss"中提取"value_loss"）
        parts = full_tag.split('/')
        if len(parts) >= 2:
            metric_type = parts[-1]
        else:
            metric_type = full_tag
        
        if metric_type not in metric_types:
            metric_types[metric_type] = []
        
        metric_types[metric_type].append((full_tag, data))
    
    # 为每种指标类型创建一个图表，包含所有智能体的数据
    for metric_type, metrics in metric_types.items():
        plt.figure(figsize=(12, 8))
        
        for full_tag, data in metrics:
            # 提取智能体名称作为图例
            parts = full_tag.split('/')
            if len(parts) >= 2:
                agent_name = parts[0]
                label = f"{agent_name}"
            else:
                label = full_tag
            
            plt.plot(data['steps'], data['values'], label=label)
        
        plt.title(f"{metric_type}")
        plt.xlabel('Steps')
        plt.ylabel('Value')
        plt.grid(True)
        plt.legend()
        
        if save_dir:
            plt.savefig(os.path.join(save_dir, f"{metric_type.replace('/', '_')}.png"))
        
        plt.close()  # 关闭图表以节省内存
    
    # 另外创建每个智能体的所有指标综合图
    agents = set()
    for full_tag in metrics_data.keys():
        parts = full_tag.split('/')
        if len(parts) >= 2:
            agents.add(parts[0])
    
    for agent in agents:
        # 为每个智能体创建一个子图布局
        agent_metrics = {k: v for k, v in metrics_data.items() if k.startswith(agent + '/')}
        if not agent_metrics:
            continue
        
        n_metrics = len(agent_metrics)
        n_cols = min(3, n_metrics)
        n_rows = (n_metrics + n_cols - 1) // n_cols
        
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4 * n_rows))
        if n_rows == 1 and n_cols == 1:
            axes = np.array([axes])
        axes = axes.flatten()
        
        for i, (full_tag, data) in enumerate(agent_metrics.items()):
            if i < len(axes):
                metric_name = full_tag.split('/')[-1]
                axes[i].plot(data['steps'], data['values'])
                axes[i].set_title(metric_name)
                axes[i].set_xlabel('Steps')
                axes[i].set_ylabel('Value')
                axes[i].grid(True)
        
        # 隐藏未使用的子图
        for i in range(len(agent_metrics), len(axes)):
            axes[i].axis('off')
        
        plt.tight_layout()
        
        if save_dir:
            plt.savefig(os.path.join(save_dir, f"{agent}_summary.png"))
        
        plt.close()

if __name__ == "__main__":
    # 分析run3目录
    run_dir = "f:/curWork/light_mappo/results/UAV_CPP/uav_cpp/rmappo/uav_cpp_training/run3"
    
    if os.path.exists(run_dir):
        analyze_run_directory(run_dir)
    else:
        print(f"目录不存在: {run_dir}")
        
        # 如果指定目录不存在，则查找所有日志目录
        log_dirs = find_tensorboard_logs()
        
        if not log_dirs:
            print("未找到任何TensorBoard日志文件，请检查结果目录")
        else:
            print(f"找到 {len(log_dirs)} 个日志目录:")
            for i, log_dir in enumerate(log_dirs):
                print(f"{i+1}. {log_dir}")
            
            # 让用户选择要分析的日志目录
            if len(log_dirs) == 1:
                selected_idx = 0
            else:
                selected_idx = int(input("请选择要分析的日志目录编号: ")) - 1
            
            # 分析选定的日志目录
            metrics_data = analyze_tensorboard_logs(log_dirs[selected_idx])
            
            # 绘制指标图表
            if metrics_data:
                save_dir = os.path.join(log_dirs[selected_idx], "analysis")
                plot_metrics(metrics_data, save_dir)
                print(f"\n图表已保存到: {save_dir}")